package com.zbkj.admin.util;

import lombok.Cleanup;

import java.io.*;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * @Description: 树操作方法工具类
 * @Author: 公众号：赵侠客
 * @Copyright: Copyright (c) 赵侠客
 * @Date: 2024-07-22 10:42
 * @Version: 1.0
 */
public class TreeUtil {

    /**
     * 创建树结构
     *
     * @param list           原始数据
     * @param rootCheck      x->x.getPId()==-1L
     * @param parentCheck    (x, y)->x.getId().equals(y.getPId())
     * @param setSubChildren MenuVo::setSubMenus
     * @param <E>            数据类型
     */
    public static <E> List<E> makeTree(List<E> list, Predicate<E> rootCheck, BiFunction<E, E, Boolean> parentCheck, BiConsumer<E, List<E>> setSubChildren) {
        return list.stream().filter(rootCheck).peek(x -> setSubChildren.accept(x, makeChildren(x, list, parentCheck, setSubChildren))).collect(Collectors.toList());
    }

    private static <E> List<E> makeChildren(E parent, List<E> allData, BiFunction<E, E, Boolean> parentCheck, BiConsumer<E, List<E>> setSubChildren) {
        return allData.stream().filter(x -> parentCheck.apply(parent, x)).peek(x -> setSubChildren.accept(x, makeChildren(x, allData, parentCheck, setSubChildren))).collect(Collectors.toList());
    }

    /**
     * 使用Map合成树
     *
     * @param menuList       需要合成树的List
     * @param pId            对象中的父ID字段,如:Menu:getPid
     * @param id             对象中的id字段 ,如：Menu:getId
     * @param rootCheck      判断E中为根节点的条件，如：x->x.getPId()==-1L , x->x.getParentId()==null,x->x.getParentMenuId()==0
     * @param setSubChildren E中设置下级数据方法，如：Menu::setSubMenus
     * @param <T>            ID字段类型
     * @param <E>            泛型实体对象
     * @return
     */
    public static <T, E> List<E> makeTree(List<E> menuList, Function<E, T> pId, Function<E, T> id, Predicate<E> rootCheck, BiConsumer<E, List<E>> setSubChildren) {
        if (menuList.size() == 1) {
            return menuList;
        }
        List<E> es;
        try {
            es = Objects.nonNull(menuList) ? deepCopy(menuList) : null;
        } catch (IOException | ClassNotFoundException e) {
            throw new RuntimeException(e);
        }
        //按原数组顺序构建父级数据Map，使用Optional考虑pId为null
        Map<Optional<T>, List<E>> parentMenuMap = es.parallelStream().collect(Collectors.groupingBy(
                node -> Optional.ofNullable(pId.apply(node)),
                LinkedHashMap::new,
                Collectors.toList()
        ));
        List<E> result = new ArrayList<>();
        for (E node : es) {
            //添加到下级数据中
            setSubChildren.accept(node, parentMenuMap.get(Optional.ofNullable(id.apply(node))));
            //如里是根节点，加入结构
            if (rootCheck.test(node)) {
                result.add(node);
            }
        }
        return result;
    }

    /**
     * 树中搜索
     *
     * @param tree
     * @param predicate
     * @param getSubChildren
     * @param <E>
     * @return 返回搜索到的节点及其父级到根节点
     * {@code List<MenuVo> searchRes = TreeUtil.search(tree1,x->"赵侠客".equals(x.getName()) ,MenuVo::getSubMenus);}
     */
    public static <E> List<E> search(List<E> tree, Predicate<E> predicate, Function<E, List<E>> getSubChildren) {
        List<E> es;
        try {
            es = Objects.nonNull(tree) ? deepCopy(tree) : null;
        } catch (IOException | ClassNotFoundException e) {
            throw new RuntimeException(e);
        }
        Iterator<E> iterator = es.iterator();
        while (iterator.hasNext()) {
            E item = iterator.next();
            List<E> childList = getSubChildren.apply(item);
            if (childList != null && !childList.isEmpty()) {
                search(childList, predicate, getSubChildren);
            }
            if (!predicate.test(item) && (childList == null || childList.isEmpty())) {
                iterator.remove();
            }
        }
        return tree;
    }

    /**
     * <pre>{@code
     * // 假设已有树结构 tree，要查找 name="目标节点" 的节点
     * List<MenuVo> result = TreeUtil.getThreeLevelNodes(
     *     tree,
     *     x -> "目标节点".equals(x.getName()),
     *     MenuVo::getSubMenus,
     *     MenuVo::setSubMenus
     * ); }</pre>
     *
     * @param tree
     * @param targetPredicate
     * @param getSubChildren
     * @param setSubChildren
     * @param <E>
     * @return
     */
    public static <E> List<E> getThreeLevelNodes(List<E> tree,
                                                 Predicate<E> targetPredicate,
                                                 Function<E, List<E>> getSubChildren,
                                                 BiConsumer<E, List<E>> setSubChildren) {
        // 1. 深拷贝原始树
        List<E> clonedTree;
        try {
            clonedTree = deepCopy(tree);
        } catch (IOException | ClassNotFoundException e) {
            throw new RuntimeException("Deep copy failed", e);
        }

        // 2. 查找目标节点及其路径
        List<E> path = new ArrayList<>();
        findNodePath(clonedTree, targetPredicate, getSubChildren, path);

        // 3. 收集前后三级节点
        return collectThreeLevelNodes(path, getSubChildren, setSubChildren);
    }

    private static <E> boolean findNodePath(List<E> nodes, Predicate<E> predicate, Function<E, List<E>> getSubChildren, List<E> path) {
        for (E node : nodes) {
            path.add(node);
            if (predicate.test(node)) {
                return true;
            }
            List<E> children = getSubChildren.apply(node);
            if (children != null && !children.isEmpty()) {
                if (findNodePath(children, predicate, getSubChildren, path)) {
                    return true;
                }
            }
            path.removeLast();
        }
        return false;
    }

    private static <E> List<E> collectThreeLevelNodes(List<E> path,
                                                      Function<E, List<E>> getSubChildren,
                                                      BiConsumer<E, List<E>> setSubChildren) {
        List<E> result = new ArrayList<>();
        if (path.isEmpty()) return result;

        // 收集向上三级父节点 Math.max(0, path.size() - 4); 前后三级 Math.max(0, path.size() - 2); 前后一级
        int startIndex = Math.max(0, path.size() - 2); // 包含自己共三级
        List<E> ancestors = path.subList(startIndex, path.size());
        result.addAll(ancestors);

        // 收集向下三级子节点
        E targetNode = path.getLast();
        collectDescendants(targetNode, 3, getSubChildren, setSubChildren, result);

        return result.stream().distinct().collect(Collectors.toList());
    }

    @SuppressWarnings("all")
    private static <E> void collectDescendants(E node, int remainingLevels, Function<E, List<E>> getSubChildren, BiConsumer<E, List<E>> setSubChildren, List<E> result) {
        if (remainingLevels <= 0) return;

        List<E> children = Optional.ofNullable(getSubChildren.apply(node)).orElseGet(ArrayList::new);

        // 临时清空子节点避免循环引用
        setSubChildren.accept(node, new ArrayList<>());

        result.addAll(children);
        children.parallelStream().peek(child -> collectDescendants(child, remainingLevels - 1, getSubChildren, setSubChildren, result)).toList();

        // 恢复原始子节点
        setSubChildren.accept(node, children);
    }


    public static <T> List<T> deepCopy(List<T> src) throws IOException, ClassNotFoundException {
        @Cleanup
        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
        ObjectOutputStream out = new ObjectOutputStream(byteOut);
        out.writeObject(src);

        @Cleanup
        ByteArrayInputStream byteIn = new ByteArrayInputStream(byteOut.toByteArray());
        ObjectInputStream in = new ObjectInputStream(byteIn);
        @SuppressWarnings("unchecked")
        List<T> dest = (List<T>) in.readObject();
        return dest;
    }

}